import torch
import numpy as np

class splitByShape:
    def __init__(self):
        pass
    
    # 0 : [Batch_Tensor,Batch_Tensor,...] -> [[Tensor,Tensor,...],[Tensor,Tensor,...],...]
    # 1 : [[Batch_Tensor,Batch_Tensor,Batch_Tensor],[Batch_Tensor,Batch_Tensor,Batch_Tensor],...] -> [[[Tensor,Tensor,Tensor],[Tensor,Tensor,Tensor],...],[[Tensor,Tensor,Tensor],[Tensor,Tensor,Tensor],...],...]
    # DONE: 完善方法使其支持上述泛型
    # 功能描述：将Batch_Tensor的shape[0]分离开进行输出。
    def process(self, data):
        res = []
        for item in data:
            items = []
            if isinstance(item, list):
                # type 1
                x = len(item)    # 3
                y = len(item[0]) # n
                for i in range(y):
                    group = []
                    for j in range(x):
                        group.append(item[j][i].copy())
                    items.append(group)
            else:
                # type 0
                for i in range(item.shape[0]):
                    items.append(item[i].copy())
            res.append(items)
        return res

